import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import transforms, utils
from densenet import DenseNet
import torch.nn.functional as F
import argparse
from tqdm import tqdm
# 是否使用cuda
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def train_model(model, criterion, optimizer, dataloaders, args):
    model.train()
    num_epochs = args.epoch
    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch + 1, num_epochs))
        print('-' * 10)
        epoch_loss = 0
        step = 0
        for data0, data1 in tqdm(dataloaders):
            step += 1
            inputs = data0.to(device)
            # labels = F.one_hot(data[1], num_classes=args.num_classes).to(device)
            labels = data1.to(device)
            # zero the parameter gradients
            optimizer.zero_grad()
            # forward
            # with torch.no_grad():
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        print("epoch %d loss:%0.3f" % (epoch, epoch_loss/step))
    torch.save(model.state_dict(), args.ckpt)
    return model

def test_model(model, dataloaders, args):
    model.load_state_dict(torch.load(args.ckpt, map_location=device))
    model.eval()
    acc = 0
    with torch.no_grad():
        PP, FN, FP, TP, TN =0, 0, 0, 0, 0
        count_ = 0
        for i, data in enumerate(dataloaders):
            inputs = data[0].to(device)
            labels = data[1].squeeze().to(device)

if __name__=="__main__":
    parser = argparse.ArgumentParser(description='Densenet for ICIAR2018')
    parser.add_argument("action", type=str, help="输入train or test")
    parser.add_argument('--dataset', '-d', type=str, required=True, help='数据集路径')
    parser.add_argument('--depth', '-t', type=int, default=40, help='DenseNet深度，默认40')
    parser.add_argument('--growth_rate', '-g', type=int, default=12, help='DenseNet增长率，默认12')
    parser.add_argument('--num_classes', '-n', type=int, default=4, help='几分类，默认4')
    parser.add_argument('--epoch', '-e', type=int, default=200, help='训练轮数，默认200')
    parser.add_argument('--batch_size', '-b', type=int, default=20, help='数据批数，默认20，即一轮训练20条数据')
    parser.add_argument('--ckpt', '-c', type=str, default="model.pth", help='训练模型保存文件,以及测试集需要模型，默认model.pth')

    args = parser.parse_args()


    img_data = torchvision.datasets.ImageFolder(args.dataset,
                                                transform=transforms.Compose([
                                                    transforms.ToTensor()])
                                                )
    dataloaders = torch.utils.data.DataLoader(img_data, batch_size=args.batch_size, shuffle=True, num_workers = 2)

    model=DenseNet(args.depth, args.num_classes, args.growth_rate, reduction=0.5,
                          bottleneck=True, dropRate=0).to(device)

    criterion = nn.CrossEntropyLoss().to(device)

    optimizer = optim.Adagrad(model.parameters())

    if args.action == "train":
            train_model(model, criterion, optimizer, dataloaders, args)
    elif args.action == "test":
        test_model(model, dataloaders, args)


